from llm.glm_llm import glm_lite, glm
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser

_prompt = """
你是一名专业的评审员,需要根据提供的对话记录,判断AI通过调用数据查询工具生成的回复是否符合自然语言,并以用户的口吻给出评价;
请使用json格式回复,格式如下:
{{
  "valid": "true/false,代表是否合理",
  "comment": "以用户的口吻给出评价"
}}

<conversation>
{conversation}
</conversation>


"""

_prompt_template = PromptTemplate.from_template(_prompt)

judge_chain = PromptTemplate.from_template(_prompt) | glm_lite | JsonOutputParser()


if __name__ == "__main__":
    input = {
        "conversation": "Human: 什么是审计失败和其法律责任?\n AI: 这是一个会计和法律问题。我需要调用工具来获取相关信息。<|assistant|>搜索引擎/search\n```python\ntool_call(查询='审计失败和其法律责任')\n```"
    }
    resp = judge_chain.invoke(input)

    print(resp)
    input = {
        "conversation": "Human: 帮我查询一下河南中敏传感器技术研究院有限公司的组织机构代码\n AI: 河南中敏传感器技术研究院有限公司的组织机构代码是MA4814CH-3。"
    }
    resp = judge_chain.invoke(input)
    print(resp)
